
import matplotlib.pyplot as plt
import os

def loss_curves(train_loss, valid_loss):

    epochs = range(len(train_loss))
    plt.figure()
    plt.plot(epochs, train_loss, 'red', linewidth = 2, label="train loss")
    plt.plot(epochs, valid_loss, 'coral', linewidth = 2, label="valid loss")
    plt.xlabel("Epoch")
    plt.ylabel("Loss")
    plt.legend(loc="upper right")  # 设置图例
    plt.grid(True)
    plt.savefig(os.path.join("logs", "epoch_loss.png"))
    plt.cla()  # 清除axes,即当前figure中的活动的axes
    plt.close("all")  # 关闭window


def lr_curves(lr):
    epochs = range(len(lr))
    plt.figure()
    plt.plot(epochs, lr, 'red', linewidth=2, label="learning rate")
    plt.xlabel("Epoch")
    plt.ylabel("lr")
    plt.legend(loc="upper right")  # 设置图例
    plt.grid(True)
    plt.savefig(os.path.join("logs", "epoch_lr.png"))
    plt.cla()  # 清除axes,即当前figure中的活动的axes
    plt.close("all")  # 关闭window


if __name__ == "__main__":
    y = [55, 2, 6, 7, 8, 235.5, 58.5, 8868, 25425] * 100
    x = [55, 2, 6, 7, 8, 55, 6, 58, 5] * 100
    loss_curves(y, x)
